package compress

import (
	"changeme/backend/utils"
	"fmt"
	"github.com/yeka/zip"
	"io"
	"io/fs"
	"os"
	"path/filepath"
	"strings"
)

type zipOption struct {
	password     string
	encrypt      zip.EncryptionMethod
	encryptFiles map[string]int
}

type ZipOptions = func(o *zipOption)

func ZipWithPassword(password string) ZipOptions {
	return func(o *zipOption) {
		o.password = password
	}
}

func ZipWithEncryptionMethod(encrypt zip.EncryptionMethod) ZipOptions {
	return func(o *zipOption) {
		o.encrypt = encrypt
	}
}

func ZipWithEncryptionFiles(encryptFiles []string) ZipOptions {
	return func(o *zipOption) {
		for _, file := range encryptFiles {
			if _, ok := o.encryptFiles[file]; !ok {
				o.encryptFiles[file] = 1
			} else {
				o.encryptFiles[file]++
			}
		}
	}
}

func Zip(path string, createPath string, zipOptions ...ZipOptions) (err error) {
	if _, exist := utils.Exist(path); !exist {
		return fs.ErrNotExist
	}

	options := &zipOption{
		encrypt:      zip.AES256Encryption,
		encryptFiles: make(map[string]int, 0),
	}
	for _, option := range zipOptions {
		option(options)
	}

	zipFile, err := os.OpenFile(createPath, os.O_WRONLY|os.O_TRUNC|os.O_CREATE, 0755)
	defer func() {
		_ = zipFile.Close()
	}()
	if err != nil {
		return
	}

	zipWriter := zip.NewWriter(zipFile)
	defer func() {
		_ = zipWriter.Close()
	}()

	return zipCompress("", path, zipWriter, options)
}

func zipCompress(prefix string, path string, w *zip.Writer, options *zipOption) (err error) {
	info, exist := utils.Exist(path)
	if !exist {
		return fs.ErrNotExist
	}

	file, err := os.Open(path)
	defer func() {
		_ = file.Close()
	}()
	if err != nil {
		return err
	}

	if info.IsDir() {
		prefix = filepath.Join(prefix, info.Name())
		fileInfos, err := file.Readdir(-1)
		if err != nil {
			return err
		}

		for _, fileInfo := range fileInfos {
			err = zipCompress(prefix, filepath.Join(file.Name(), fileInfo.Name()), w, options)
			if err != nil {
				return err
			}
		}
	} else {
		var fw io.Writer
		header, err := zip.FileInfoHeader(info)
		if err != nil {
			return err
		}
		header.Name = filepath.Join(prefix, header.Name)

		if len(options.password) == 0 {
			fw, err = w.CreateHeader(header)
		} else {
			if _, ok := options.encryptFiles[header.Name]; len(options.encryptFiles) == 0 || ok {
				fw, err = w.Encrypt(header.Name, options.password, options.encrypt)
			} else {
				fw, err = w.CreateHeader(header)
			}
		}

		if err != nil {
			return err
		}

		_, err = io.Copy(fw, file)
		if err != nil {
			return err
		}
	}
	return
}

func Unzip(absFilePath string, extractPath string, zipOptions ...ZipOptions) (err error) {
	if !utils.ExistFile(absFilePath) {
		return fs.ErrNotExist
	}

	options := &zipOption{
		encrypt:      zip.AES256Encryption,
		encryptFiles: make(map[string]int, 0),
	}
	for _, option := range zipOptions {
		option(options)
	}

	zipReader, _ := zip.OpenReader(absFilePath)
	defer func() {
		_ = zipReader.Close()
	}()

	path, err := filepath.Abs(extractPath)
	if err != nil {
		return err
	}

	fmt.Printf("解压 %s \n\t=> %s\n", absFilePath, path)
	total := len(zipReader.File)
	var schedule int

	defer func() {
		if schedule == 50 {
			fmt.Println()
		}
		if err != nil {
			fmt.Println()
		}
	}()
	for num, file := range zipReader.File {
		printStr := "["
		schedule = int(float64(num+1) / float64(total) * 50)
		for i := 0; i <= schedule; i++ {
			printStr += "="
		}
		for j := 0; j < 50-schedule; j++ {
			printStr += " "
		}
		printStr += "]"

		fmt.Printf("\r%s %d/%d", printStr, num+1, total)

		if file.IsEncrypted() {
			file.SetPassword(options.password)
		}

		err = os.MkdirAll(filepath.Join(path, utils.SubString(file.Name, 0, strings.LastIndex(file.Name, string(filepath.Separator)))), 0755)
		if err != nil {
			return err
		}
		src, err := file.Open()
		if err != nil {
			return err
		}
		target, err := os.Create(filepath.Join(path, file.Name))
		if err != nil {
			return err
		}
		_, err = io.Copy(target, src)
		if err != nil {
			return err
		}
		_ = target.Close()
		_ = src.Close()
	}
	return
}
